# Copyright 2025 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""expected output of activation"""
import numpy as np
import mindspore as ms
from mindformers.parallel_core.training_graph.activation import GELU, SiLU, SwiGlu


def get_input(seq_length=8, batch_size=1, hidden_size=16):
    """get input_tensor for test"""
    np.random.seed(1)
    input_numpy = np.random.rand(seq_length, batch_size, hidden_size)
    input_tensor = ms.tensor(input_numpy / 100, dtype=ms.float32)
    return input_tensor


def get_output(activation):
    """output for all cases"""
    if isinstance(activation, GELU):
        gpu_output = np.array(
            [[[0.0020904541, 0.0036163330, 0.0000005737, 0.0015106201,
               0.0007324219, 0.0004615784, 0.0009307861, 0.0017318726,
               0.0019836426, 0.0027160645, 0.0020904541, 0.0034484863,
               0.0010223389, 0.0044250488, 0.0001373291, 0.0033721924]],
             [[0.0020904541, 0.0028076172, 0.0007019043, 0.0009918213,
               0.0040283203, 0.0048828125, 0.0015716553, 0.0034790039,
               0.0044250488, 0.0045166016, 0.0004253387, 0.0001955032,
               0.0008506775, 0.0044250488, 0.0004920959, 0.0021057129]],
             [[0.0048217773, 0.0026855469, 0.0034790039, 0.0015869141,
               0.0034484863, 0.0042114258, 0.0000915527, 0.0037689209,
               0.0049743652, 0.0037536621, 0.0014038086, 0.0039672852,
               0.0005149841, 0.0022583008, 0.0045776367, 0.0014648438]],
             [[0.0014419556, 0.0006484985, 0.0000967979, 0.0034027100,
               0.0010604858, 0.0013275146, 0.0024719238, 0.0002670288,
               0.0028839111, 0.0007324219, 0.0029602051, 0.0035095215,
               0.0005111694, 0.0020751953, 0.0034942627, 0.0020751953]],
             [[0.0002498627, 0.0027008057, 0.0033416748, 0.0025939941,
               0.0047607422, 0.0029449463, 0.0045471191, 0.0006866455,
               0.0006980896, 0.0040588379, 0.0019836426, 0.0008277893,
               0.0046691895, 0.0017471313, 0.0037689209, 0.0036468506]],
             [[0.0044555664, 0.0031280518, 0.0037689209, 0.0017547607,
               0.0013504028, 0.0045166016, 0.0021362305, 0.0048522949,
               0.0033264160, 0.0031280518, 0.0005722046, 0.0047912598,
               0.0022583008, 0.0029144287, 0.0020446777, 0.0011825562]],
             [[0.0045471191, 0.0028839111, 0.0000143647, 0.0030975342,
               0.0016403198, 0.0026550293, 0.0044555664, 0.0017929077,
               0.0045776367, 0.0031280518, 0.0000791550, 0.0046691895,
               0.0034637451, 0.0050048828, 0.0008621216, 0.0006866455]],
             [[0.0046997070, 0.0034942627, 0.0003299713, 0.0037994385,
               0.0037841797, 0.0046386719, 0.0035705566, 0.0006217957,
               0.0000991821, 0.0001306534, 0.0001411438, 0.0012283325,
               0.0043334961, 0.0027160645, 0.0027770996, 0.0042419434]]])
        golden_output = np.array(
            [[[0.0020920476, 0.0036223219, 0.0000005719, 0.0015153097,
               0.0007346388, 0.0004620331, 0.0009326850, 0.0017325674,
               0.0019901181, 0.0027056660, 0.0021029832, 0.0034448293,
               0.0010239289, 0.0044213491, 0.0001369679, 0.0033702708]],
             [[0.0020934714, 0.0028059017, 0.0007027210, 0.0009920727,
               0.0040293029, 0.0048787105, 0.0015710397, 0.0034807350,
               0.0044125873, 0.0045049619, 0.0004255096, 0.0001953348,
               0.0008503026, 0.0044214767, 0.0004921200, 0.0021126126]],
             [[0.0048260521, 0.0026771664, 0.0034784826, 0.0015815498,
               0.0034513066, 0.0042009177, 0.0000914547, 0.0037731696,
               0.0049833152, 0.0037631593, 0.0014053574, 0.0039712493,
               0.0005165552, 0.0022474704, 0.0045759114, 0.0014715098]],
             [[0.0014421808, 0.0006508174, 0.0000968497, 0.0034125622,
               0.0010599272, 0.0013305466, 0.0024675061, 0.0002669263,
               0.0028837374, 0.0007345017, 0.0029603823, 0.0035183257,
               0.0005120900, 0.0020771201, 0.0034912371, 0.0020777402]],
             [[0.0002498669, 0.0026909390, 0.0033365516, 0.0025850215,
               0.0047585694, 0.0029465014, 0.0045495681, 0.0006881273,
               0.0006971554, 0.0040629623, 0.0019946930, 0.0008278618,
               0.0046718628, 0.0017436543, 0.0037765487, 0.0036510169]],
             [[0.0044476567, 0.0031338786, 0.0037772090, 0.0017493485,
               0.0013525459, 0.0045114495, 0.0021477668, 0.0048613367,
               0.0033347672, 0.0031238971, 0.0005742552, 0.0047834110,
               0.0022576356, 0.0029052941, 0.0020473297, 0.0011873763]],
             [[0.0045494554, 0.0028815262, 0.0000143520, 0.0031009191,
               0.0016374809, 0.0026463734, 0.0044610221, 0.0017914409,
               0.0045756055, 0.0031323028, 0.0000791162, 0.0046816478,
               0.0034735270, 0.0050262962, 0.0008628874, 0.0006864290]],
             [[0.0046976754, 0.0035034609, 0.0003301746, 0.0038000841,
               0.0037920545, 0.0046491111, 0.0035778205, 0.0006219710,
               0.0000994164, 0.0001310823, 0.0001415644, 0.0012334738,
               0.0043296483, 0.0027057379, 0.0027763024, 0.0042384397]]])
    if isinstance(activation, SiLU):
        gpu_output = np.array(
            [[[0.0020904541, 0.0036163330, 0.0000005737, 0.0015106201,
               0.0007324219, 0.0004615784, 0.0009307861, 0.0017242432,
               0.0019836426, 0.0027008057, 0.0020904541, 0.0034484863,
               0.0010223389, 0.0044250488, 0.0001373291, 0.0033721924]],
             [[0.0020904541, 0.0028076172, 0.0007019043, 0.0009918213,
               0.0040283203, 0.0048828125, 0.0015640259, 0.0034790039,
               0.0044250488, 0.0045166016, 0.0004253387, 0.0001955032,
               0.0008506775, 0.0044250488, 0.0004920959, 0.0021057129]],
             [[0.0048217773, 0.0026702881, 0.0034790039, 0.0015792847,
               0.0034484863, 0.0042114258, 0.0000915527, 0.0037689209,
               0.0049743652, 0.0037536621, 0.0014038086, 0.0039672852,
               0.0005149841, 0.0022430420, 0.0045776367, 0.0014648438]],
             [[0.0014419556, 0.0006484985, 0.0000967979, 0.0034027100,
               0.0010604858, 0.0013275146, 0.0024566650, 0.0002670288,
               0.0028839111, 0.0007324219, 0.0029602051, 0.0035095215,
               0.0005111694, 0.0020751953, 0.0034942627, 0.0020751953]],
             [[0.0002498627, 0.0026855469, 0.0033416748, 0.0025787354,
               0.0047607422, 0.0029449463, 0.0045471191, 0.0006866455,
               0.0006980896, 0.0040588379, 0.0019836426, 0.0008277893,
               0.0046691895, 0.0017395020, 0.0037689209, 0.0036468506]],
             [[0.0044555664, 0.0031280518, 0.0037689209, 0.0017471313,
               0.0013504028, 0.0045166016, 0.0021362305, 0.0048522949,
               0.0033264160, 0.0031280518, 0.0005722046, 0.0047912598,
               0.0022430420, 0.0029144287, 0.0020446777, 0.0011825562]],
             [[0.0045471191, 0.0028839111, 0.0000143647, 0.0030975342,
               0.0016326904, 0.0026397705, 0.0044555664, 0.0017852783,
               0.0045776367, 0.0031280518, 0.0000791550, 0.0046691895,
               0.0034637451, 0.0050048828, 0.0008621216, 0.0006866455]],
             [[0.0046997070, 0.0034942627, 0.0003299713, 0.0037994385,
               0.0037841797, 0.0046386719, 0.0035705566, 0.0006217957,
               0.0000991821, 0.0001306534, 0.0001411438, 0.0012283325,
               0.0043334961, 0.0027008057, 0.0027618408, 0.0042419434]]])
        golden_output = np.array(
            [[[0.0020894576, 0.0036145942, 0.0000005719, 0.0015139480,
               0.0007343179, 0.0004619061, 0.0009321683, 0.0017307890,
               0.0019877732, 0.0027013416, 0.0021003657, 0.0034378357,
               0.0010233063, 0.0044098645, 0.0001369567, 0.0033635758]],
             [[0.0020908774, 0.0028012525, 0.0007024274, 0.0009914886,
               0.0040197521, 0.0048647462, 0.0015695768, 0.0034735959,
               0.0044011474, 0.0044930410, 0.0004254019, 0.0001953121,
               0.0008498731, 0.0044099907, 0.0004919760, 0.0021099714]],
             [[0.0048123864, 0.0026729330, 0.0034713529, 0.0015800669,
               0.0034442865, 0.0041905432, 0.0000914497, 0.0037647893,
               0.0049687512, 0.0037548218, 0.0014041861, 0.0039619706,
               0.0005163964, 0.0022444827, 0.0045636161, 0.0014702260]],
             [[0.0014409472, 0.0006505655, 0.0000968442, 0.0034056981,
               0.0010592602, 0.0013294962, 0.0024639068, 0.0002668839,
               0.0028788282, 0.0007341811, 0.0029552097, 0.0035110333,
               0.0005119340, 0.0020745662, 0.0034840554, 0.0020751848]],
             [[0.0002498297, 0.0026866617, 0.0033299886, 0.0025810732,
               0.0047452804, 0.0029413765, 0.0045374129, 0.0006878460,
               0.0006968667, 0.0040532532, 0.0019923379, 0.0008274546,
               0.0046590501, 0.0017418528, 0.0037681535, 0.0036431665]],
             [[0.0044360361, 0.0031280851, 0.0037688101, 0.0017475350,
               0.0013514609, 0.0044994960, 0.0021450375, 0.0048474730,
               0.0033282116, 0.0031181413, 0.0005740590, 0.0047699842,
               0.0022546209, 0.0029003113, 0.0020448484, 0.0011865394]],
             [[0.0045373002, 0.0028766252, 0.0000143518, 0.0030952462,
               0.0016358920, 0.0026422353, 0.0044493326, 0.0017895397,
               0.0045633116, 0.0031265151, 0.0000791125, 0.0046687825,
               0.0034664180, 0.0050114808, 0.0008624450, 0.0006861489]],
             [[0.0046847207, 0.0034962296, 0.0003301097, 0.0037915835,
               0.0037835892, 0.0046364213, 0.0035702805, 0.0006217409,
               0.0000994106, 0.0001310721, 0.0001415525, 0.0012325709,
               0.0043186308, 0.0027014138, 0.0027717503, 0.0042278799]]])
    if isinstance(activation, SwiGlu):
        gpu_output = np.array(
            [[[0.0000082850, 0.0000195503, 0.0000000024, 0.0000103712,
               0.0000014976, 0.0000040531, 0.0000002552, 0.0000115633]],
             [[0.0000183582, 0.0000251532, 0.0000005960, 0.0000003874,
               0.0000068545, 0.0000429153, 0.0000015423, 0.0000146627]],
             [[0.0000476837, 0.0000199080, 0.0000097752, 0.0000124574,
               0.0000035465, 0.0000188351, 0.0000008345, 0.0000110269]],
             [[0.0000082850, 0.0000009499, 0.0000005700, 0.0000237226,
               0.0000010878, 0.0000055134, 0.0000170469, 0.0000011101]],
             [[0.0000003483, 0.0000216961, 0.0000132322, 0.0000042617,
               0.0000441074, 0.0000102520, 0.0000340939, 0.0000049770]],
             [[0.0000295639, 0.0000194311, 0.0000043213, 0.0000166893,
               0.0000060499, 0.0000262260, 0.0000087619, 0.0000115037]],
             [[0.0000412464, 0.0000180006, 0.0000000023, 0.0000287294,
               0.0000112653, 0.0000262260, 0.0000076890, 0.0000024587]],
             [[0.0000009313, 0.0000009127, 0.0000000931, 0.0000093579,
               0.0000326633, 0.0000250340, 0.0000196695, 0.0000052452]]])
        golden_output = np.array(
            [[[0.0000082903, 0.0000194760, 0.0000000024, 0.0000103739,
               0.0000015013, 0.0000040561, 0.0000002553, 0.0000116044]],
             [[0.0000183242, 0.0000250602, 0.0000005974, 0.0000003872,
               0.0000068268, 0.0000427194, 0.0000015436, 0.0000146276]],
             [[0.0000475878, 0.0000199980, 0.0000097352, 0.0000124711,
               0.0000035554, 0.0000187692, 0.0000008309, 0.0000110540]],
             [[0.0000082727, 0.0000009546, 0.0000005707, 0.0000238317,
               0.0000010840, 0.0000055049, 0.0000171094, 0.0000011054]],
             [[0.0000003480, 0.0000216919, 0.0000132426, 0.0000042679,
               0.0000440129, 0.0000102291, 0.0000340674, 0.0000049937]],
             [[0.0000294305, 0.0000194472, 0.0000043246, 0.0000165927,
               0.0000060804, 0.0000260246, 0.0000087547, 0.0000114898]],
             [[0.0000412230, 0.0000179317, 0.0000000023, 0.0000287684,
               0.0000113023, 0.0000263516, 0.0000076680, 0.0000024541]],
             [[0.0000009313, 0.0000009164, 0.0000000934, 0.0000093353,
               0.0000325399, 0.0000249825, 0.0000197373, 0.0000052353]]])
    return gpu_output, golden_output
